//-*****************************************************************************
//
// Copyright (c) 2013,
//  Sony Pictures Imageworks Inc. and
//  Industrial Light & Magic, a division of Lucasfilm Entertainment Company Ltd.
//
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:
// *       Redistributions of source code must retain the above copyright
// notice, this list of conditions and the following disclaimer.
// *       Redistributions in binary form must reproduce the above
// copyright notice, this list of conditions and the following disclaimer
// in the documentation and/or other materials provided with the
// distribution.
// *       Neither the name of Sony Pictures Imageworks, nor
// Industrial Light & Magic, nor the names of their contributors may be used
// to endorse or promote products derived from this software without specific
// prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
// "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
// LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
// A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
// OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
// LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
// DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
// THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
//
//-*****************************************************************************

#include "AlembicReader.h"
#include "TreeEntry.h"
#include "OptionsData.h"
#include "MriGeoReaderPlugin.h"
#include <cmath>
#include <Alembic/AbcCoreFactory/IFactory.h>

using namespace Alembic;
using namespace Alembic::Abc;
using namespace std;

namespace mariAlembic
{

bool AlembicReader::loadArchive(const std::string & path)
{
    freeArchive();

    try
    {
        Alembic::AbcCoreFactory::IFactory factory;
        archive = factory.getArchive(path);
    }
    catch (const std::exception & e)
    {
        std::cerr << e.what() << std::endl;
        return false;
    }

    if (archive.valid())
    {
        return true;
    }

    std::cerr << "Error: invalid Alembic archive: " << path << std::endl;
    return false;

}


void AlembicReader::freeArchive(void)
{
    if (archive.valid())
    {
        archive.reset();
    }
}


int AlembicReader::getType(Alembic::Abc::IObject & iObj)
{
    if ( Alembic::AbcGeom::IXform::matches(iObj.getHeader()) )
        return ITEM_TRANSFORM;
    else if ( Alembic::AbcGeom::ISubD::matches(iObj.getHeader()) )
        return ITEM_SUBD;
    else if ( Alembic::AbcGeom::IPolyMesh::matches(iObj.getHeader()) )
        return ITEM_POLYMESH;
    else if ( Alembic::AbcGeom::ICamera::matches(iObj.getHeader()) )
        return ITEM_CAMERA;
    else if ( Alembic::AbcGeom::ICurves::matches(iObj.getHeader()) )
        return ITEM_CURVES;
    else if ( Alembic::AbcGeom::INuPatch::matches(iObj.getHeader()) )
        return ITEM_NURBSPATCH;
    else if ( Alembic::AbcGeom::IPoints::matches(iObj.getHeader()) )
        return ITEM_POINTS;

    return ITEM_UNKNOWN;
}


string AlembicReader::getShortName(string &path)
{
    size_t pos;
    pos = path.rfind('/');
    if (pos != string::npos)
    {
        string p2 = &path[pos+1];
        return p2;
    }
    return path;
}


int AlembicReader::retrieveObjectData(IObject iObj, vector<string> &shapeNames,
    TreeEntry &tree)
{
    TreeEntry *tentry = new TreeEntry;
    tentry->fullName = iObj.getFullName();
    tentry->shortName = getShortName(tentry->fullName);

    // get item type
    int itemType = getType(iObj);
    if (itemType == ITEM_POLYMESH)
    {
        tentry->type = ITEM_POLYMESH;

        // it's a polymesh, so add it to shapeNames
        shapeNames.push_back(tentry->fullName);

        // get arbitrary parameters
        AbcGeom::IPolyMesh meshObject(iObj, AbcGeom::kWrapExisting);
        Abc::ICompoundProperty arbGeomParams =
            meshObject.getSchema().getArbGeomParams();
        processArbitraryGeomParams(arbGeomParams, meshObject, tentry->fullName);
    }
    else
    {
        tentry->type = ITEM_UNKNOWN;
    }

    tree.children.push_back(tentry);

    for ( unsigned int i = 0 ; i < iObj.getNumChildren() ; i++ )
    {
        retrieveObjectData(IObject(iObj, iObj.getChildHeader(i).getName()),
            shapeNames, *tentry);
    }

    return 1;
}


bool AlembicReader::queryAlembicData(vector<string> &shapeNames)
{

    if(loadArchive(pFileName))
    {
        shapeNames.clear();
        if (archive.valid())
        {
            root = archive.getTop();

            // start with the children of the root node
            for ( unsigned int i = 0 ; i < root.getNumChildren() ; i++ )
            {
                retrieveObjectData(IObject(root,
                    root.getChildHeader(i).getName()),
                    shapeNames, widgetData->rootTree);
            }
        }
        else
        {
            return false;
        }

        freeArchive();
        return true;
    }

    return false;
}


M44d AlembicReader::processTransform(IObject &iObj)
{
    M44d curMatrix;
    curMatrix.makeIdentity();

    Alembic::AbcGeom::IXform xform(iObj, Alembic::Abc::kWrapExisting);

    Alembic::AbcGeom::XformSample samp;
    xform.getSchema().get(samp, 0);

    M44d tmpMat;
    double r;

    int numOps = samp.getNumOps();
    for (int i = (numOps-1); i >= 0;  i--) // perform ops in reverse order
    {
        AbcGeom::XformOp xop = samp.getOp(i);
        switch(xop.getType())
        {
            case  AbcGeom::kScaleOperation:
                tmpMat.makeIdentity();
                tmpMat.setScale(xop.getScale());
                curMatrix *= tmpMat;
                break;
            case  AbcGeom::kTranslateOperation:
                tmpMat.makeIdentity();
                tmpMat.setTranslation(xop.getTranslate());
                curMatrix *= tmpMat;
                break;
            case  AbcGeom::kRotateOperation:
                break;
            case  AbcGeom::kMatrixOperation:
                tmpMat.makeIdentity();
                tmpMat = xop.getMatrix();
                curMatrix *= tmpMat;
                break;
            case  AbcGeom::kRotateXOperation:
                r = (xop.getXRotation() * DTOR);
                tmpMat.makeIdentity();
                tmpMat = M44d(1, 0, 0, 0,
                              0, cos(r), sin(r), 0,
                              0, -(sin(r)),cos(r), 0,
                              0, 0, 0, 1);
                curMatrix *= tmpMat;
                break;
            case  AbcGeom::kRotateYOperation:
                r = (xop.getYRotation() * DTOR);
                tmpMat.makeIdentity();
                tmpMat = M44d(cos(r), 0, -(sin(r)), 0,
                              0, 1, 0, 0,
                              sin(r), 0, cos(r), 0,
                              0, 0, 0, 1);
                curMatrix *= tmpMat;
                break;
            case  AbcGeom::kRotateZOperation:
                r = (xop.getZRotation() * DTOR);
                tmpMat.makeIdentity();
                tmpMat = M44d(cos(r), sin(r), 0, 0,
                              -(sin(r)), cos(r), 0, 0,
                              0, 0, 1, 0,
                              0, 0, 0, 1);
                curMatrix *= tmpMat;
                break;
        }
    }

    return curMatrix;
}


void AlembicReader::getNormals(AbcGeom::IPolyMesh &meshObject,
                               Alembic::Abc::N3fArraySamplePtr &normalVal,
                               size_t *normalSize, string &path)
{
    AbcGeom::IN3fGeomParam iNormals = meshObject.getSchema().getNormalsParam();
    if (iNormals.valid())
    {
        if (iNormals.getScope() != Alembic::AbcGeom::kVertexScope &&
            iNormals.getScope() != Alembic::AbcGeom::kVaryingScope &&
            iNormals.getScope() != Alembic::AbcGeom::kFacevaryingScope)
        {
            cout << "Normals have an unsupported scope, skipping normals for: "
                << path << endl;
        }
        else
        {
            Alembic::AbcCoreAbstract::index_t index=0;

            Alembic::AbcGeom::IN3fGeomParam::Sample samp;
            iNormals.getExpanded(samp, Alembic::Abc::ISampleSelector(index));

            normalVal = samp.getVals();
            *normalSize = normalVal->size();
        }
    }
    else
        cout << "Missing Normal parameter: '" << path <<"'" << endl;

}


void AlembicReader::getUVs(AbcGeom::IPolyMesh &meshObject,
    AbcGeom::IV2fGeomParam &uvParam,
    Alembic::AbcGeom::V2fArraySamplePtr &uvPtr,
    Alembic::Abc::UInt32ArraySamplePtr &uvIndexPtr, string &path)
{
    uvParam = meshObject.getSchema().getUVsParam();
    if (uvParam.valid())
    {
        Alembic::AbcCoreAbstract::index_t index=0;

        Alembic::AbcGeom::IV2fGeomParam::Sample samp;
        uvParam.getIndexed(samp, Alembic::Abc::ISampleSelector(index));

        uvPtr = samp.getVals();
        uvIndexPtr = samp.getIndices();
    }
    else
        cout << "Missing UV parameter: '" << path << "'" << endl;
}


bool AlembicReader::createMesh(int triVertCnt, vector<float>&tris,
    vector<int>&indices, vector<float>&normals, vector<float>&uvs,
    MriGeoEntityHandle Entity, string &path)
{

    MriGeoDataHandle Vertices, Indices, Normals, UVs;
    MriGeoObjectHandle Object;//, ChildObject;
    MriGeoPluginResult Result;

    // Create geometry data
    CHECK_RESULT(s_Host->createGeoData(Entity, tris.data(),
        tris.size() * sizeof(float), MRI_GDT_FLOAT_BUFFER, MRI_GDR_VERTICES,
        &Vertices));

    if (normals.size())
    {
        CHECK_RESULT(s_Host->createGeoData(Entity, normals.data(),
            normals.size() * sizeof(float), MRI_GDT_FLOAT_BUFFER,
            MRI_GDR_NORMALS, &Normals));
    }
    CHECK_RESULT(s_Host->createGeoData(Entity, indices.data(),
        indices.size() * sizeof(unsigned int), MRI_GDT_U32_BUFFER,
        MRI_GDR_INDICES, &Indices));

    CHECK_RESULT(s_Host->createGeoData(Entity, uvs.data(),
        uvs.size() * sizeof(float), MRI_GDT_FLOAT_BUFFER, MRI_GDR_UV, &UVs));


    // Create geo object and add data

    CHECK_RESULT(s_Host->createGeoObject(Entity, path.c_str(), (triVertCnt/3),
                 &Object));

    CHECK_RESULT(s_Host->addGeoDataToObject(Entity, Object, Vertices));
    if (normals.size())
    {
        CHECK_RESULT(s_Host->addGeoDataToObject(Entity, Object, Normals));
    }
    CHECK_RESULT(s_Host->addGeoDataToObject(Entity, Object, Indices));
    CHECK_RESULT(s_Host->addGeoDataToObject(Entity, Object, UVs));

    return true;
}


void AlembicReader::getArbStringParam(Abc::ICompoundProperty & parent,
    const std::string &propName, Abc::chrono_t sampleTime, string &val)
{

    AbcGeom::IStringGeomParam param(parent, propName);
    AbcGeom::IStringGeomParam::Sample paramSample;
    param.getExpanded(paramSample, Abc::ISampleSelector(sampleTime));
    Abc::StringArraySamplePtr paramVals = paramSample.getVals();
    if (paramVals->size())
    {
        val = (*paramVals)[0];
    }
}


void AlembicReader::getArbBoolParam(Abc::ICompoundProperty & parent,
    const std::string &propName, Abc::chrono_t sampleTime, bool &val)
{

    AbcGeom::IBoolGeomParam param(parent, propName);
    AbcGeom::IBoolGeomParam::Sample paramSample;
    param.getExpanded(paramSample, Abc::ISampleSelector(sampleTime));
    Abc::BoolArraySamplePtr paramVals = paramSample.getVals();
    if (paramVals->size())
    {
        val = (*paramVals)[0];
    }
}


void AlembicReader::getArbIntParam(Abc::ICompoundProperty & parent,
    const std::string &propName, Abc::chrono_t sampleTime, int &val)
{

    AbcGeom::IInt32GeomParam param(parent, propName);
    AbcGeom::IInt32GeomParam::Sample paramSample;
    param.getExpanded(paramSample, Abc::ISampleSelector(sampleTime));
    Abc::Int32ArraySamplePtr paramVals = paramSample.getVals();
    if (paramVals->size())
    {
        val = (*paramVals)[0];
    }
}


void AlembicReader::processArbitraryGeomParams(Abc::ICompoundProperty & parent,
    AbcGeom::IPolyMesh &meshObject, string &fullpath)
{
    if (!parent.valid())
    {
        return;
    }

    Abc::TimeSamplingPtr ts = meshObject.getSchema().getTimeSampling();

    Abc::chrono_t startTime = ts->getSampleTime(0);

    for (size_t i = 0; i < parent.getNumProperties(); ++i)
    {
        const Abc::PropertyHeader &propHeader = parent.getPropertyHeader(i);
        const std::string &propName = propHeader.getName();

        if(AbcGeom::IStringGeomParam::matches(propHeader) &&
            propName == materialAttrName)
        {
            string paramData;
            getArbStringParam(parent, propName, startTime, paramData);

            map<string, set<string> >::iterator it;
            it = widgetData->materialMap.find(paramData);
            if (it == widgetData->materialMap.end())
            {
                set<string> ss;
                ss.insert(fullpath);
                widgetData->materialMap[paramData] = ss;
            }
            else
            {
                set<string> &ss = (*it).second;
                ss.insert(fullpath);
            }

            map<string, string>::iterator sit;
            sit = widgetData->shapeMap.find(fullpath);
            if (sit == widgetData->shapeMap.end())
            {
                widgetData->shapeMap[fullpath] = paramData;
            }

        }
    }
}


void triangleFan(int numVerts, int &vIndex, const float *points,
    Alembic::Abc::Int32ArraySamplePtr &faceIndices, vector<int> &indices,
    int &indexCnt, vector<float> &tris, vector<float> &normals,
    Alembic::Abc::N3fArraySamplePtr &normalVal, int normalSize,
    vector<float> &uvs, AbcGeom::IV2fGeomParam &uvParam,
    Alembic::AbcGeom::V2fArraySamplePtr &uvPtr,
    Alembic::Abc::UInt32ArraySamplePtr &uvIndexPtr, M44d &matrix)
{

    Imath::Vec3<float> p, np;

    for (int k=0; k < (numVerts-2); k++)
    {
        p = Imath::Vec3<float>(points[(*faceIndices)[vIndex]*3],
            points[((*faceIndices)[vIndex]*3)+1],
            points[((*faceIndices)[vIndex]*3)+2]);
        matrix.multVecMatrix(p, np);
        tris.push_back(np[0]);
        tris.push_back(np[1]);
        tris.push_back(np[2]);

        p = Imath::Vec3<float>(points[(*faceIndices)[vIndex+k+1]*3],
            points[((*faceIndices)[vIndex+k+1]*3)+1],
            points[((*faceIndices)[vIndex+k+1]*3)+2]);
        matrix.multVecMatrix(p, np);
        tris.push_back(np[0]);
        tris.push_back(np[1]);
        tris.push_back(np[2]);

        p = Imath::Vec3<float>(points[(*faceIndices)[vIndex+k+2]*3],
            points[((*faceIndices)[vIndex+k+2]*3)+1],
            points[((*faceIndices)[vIndex+k+2]*3)+2]);
        matrix.multVecMatrix(p, np);
        tris.push_back(np[0]);
        tris.push_back(np[1]);
        tris.push_back(np[2]);


        indices.push_back(indexCnt++);
        indices.push_back(indexCnt++);
        indices.push_back(indexCnt++);

        // add normals
        if (normalSize)
        {
            normals.push_back((*normalVal)[vIndex].x);
            normals.push_back((*normalVal)[vIndex].y);
            normals.push_back((*normalVal)[vIndex].z);

            normals.push_back((*normalVal)[vIndex+k+1].x);
            normals.push_back((*normalVal)[vIndex+k+1].y);
            normals.push_back((*normalVal)[vIndex+k+1].z);

            normals.push_back((*normalVal)[vIndex+k+2].x);
            normals.push_back((*normalVal)[vIndex+k+2].y);
            normals.push_back((*normalVal)[vIndex+k+2].z);
        }


        // add uvs
        if (uvParam.valid() && uvPtr->size() && uvIndexPtr->size())
        {
            uvs.push_back( (*uvPtr)[(*uvIndexPtr)[vIndex]][0]);
            uvs.push_back( (*uvPtr)[(*uvIndexPtr)[vIndex]][1]);

            uvs.push_back( (*uvPtr)[(*uvIndexPtr)[vIndex+k+1]][0]);
            uvs.push_back( (*uvPtr)[(*uvIndexPtr)[vIndex+k+1]][1]);

            uvs.push_back( (*uvPtr)[(*uvIndexPtr)[vIndex+k+2]][0]);
            uvs.push_back( (*uvPtr)[(*uvIndexPtr)[vIndex+k+2]][1]);
        }
        else
        {
            // uvs are missing, add default values so the object will load
            for (int n=0; n<6; n++)
            {
                uvs.push_back(0);
            }
        }

    }

}


bool AlembicReader::processMesh(IObject &iObj, M44d matrix,
    MriGeoEntityHandle Entity, string &path)
{

    AbcGeom::IPolyMesh meshObject(iObj, AbcGeom::kWrapExisting);

    Abc::TimeSamplingPtr ts = meshObject.getSchema().getTimeSampling();

    double startTime = ts->getSampleTime(0);

    Abc::chrono_t inputSampleTime = (startTime);
    Abc::ISampleSelector sampleSelector(inputSampleTime);

    AbcGeom::IPolyMeshSchema::Sample sample =
            meshObject.getSchema().getValue(sampleSelector);

    // get points
    const float *points = reinterpret_cast<const float*>(
            sample.getPositions()->get());

    // get face counts
    unsigned int numPolys = sample.getFaceCounts()->size();
    Abc::Int32ArraySamplePtr faceCounts = sample.getFaceCounts();

    // get face indices
    unsigned int numFaceIndices = sample.getFaceIndices()->size();
    Abc::Int32ArraySamplePtr faceIndices = sample.getFaceIndices();

    // get normals
    vector<float> normals;
    Alembic::Abc::N3fArraySamplePtr normalVal;
    size_t normalSize = 0;
    getNormals(meshObject, normalVal, &normalSize, path);

    // get uvs
    AbcGeom::IV2fGeomParam uvParam;
    Alembic::AbcGeom::V2fArraySamplePtr uvPtr;
    Alembic::Abc::UInt32ArraySamplePtr uvIndexPtr;
    getUVs(meshObject, uvParam, uvPtr, uvIndexPtr, path);

    // check to see if data matches, fail if not
    if (normalSize && (numFaceIndices != normalSize))
    {
        cout << "Normals indices size (" << normalSize <<
            ") does not match face indices size (" <<
            numFaceIndices << ")" << endl;
        return false;
    }

    if(uvIndexPtr && (uvIndexPtr->size()!= numFaceIndices))
    {
        cout << "UV indices size (" << uvIndexPtr->size() <<
            ") does not match face indices size (" << numFaceIndices <<
            ")" << endl;
        return false;
    }

    int triVertCnt=0;
    for(unsigned int i=0; i < numPolys; i++)
    {
        if ((*faceCounts)[i] > 2)
        {
            triVertCnt += (((*faceCounts)[i]-2)*3);
        }
    }


    // create data to pass to Mari

    int cnt = 0;
    int vIndex = 0;

    vector<float> tris;
    tris.reserve(triVertCnt*3);

    vector<int>indices;
    indices.reserve(triVertCnt);

    vector<float> uvs;
    uvs.reserve(triVertCnt*2);

    Imath::Vec3<float> p, np;

    for(unsigned int j=0; j < numPolys; j++)
    {
        if ((*faceCounts)[j] >= 3)
        {
            triangleFan((*faceCounts)[j], vIndex, points, faceIndices, indices,
                cnt, tris, normals, normalVal, normalSize, uvs, uvParam, uvPtr,
                uvIndexPtr, matrix);
        }
        vIndex+=(*faceCounts)[j];
    }

    // pass the geometry to Mari
    int result = createMesh(triVertCnt, tris, indices, normals, uvs, Entity,
        path);
    if (result != MRI_GPR_SUCCEEDED)
        return false;

    return true;
}


int AlembicReader::processObject(IObject iObj, M44d matrix,
    MriGeoEntityHandle Entity, TreeEntry &tree)
{
    string fullpath = iObj.getFullName();
    string path = getShortName(fullpath);

    M44d curMatrix;
    curMatrix.makeIdentity();

    TreeEntry *tentry = new TreeEntry;
    tentry->fullName = fullpath;
    tentry->shortName = path;

    int itemType = getType(iObj);
    if (itemType == ITEM_TRANSFORM)
    {
        curMatrix = processTransform(iObj);
        tentry->matrix = curMatrix;
        tentry->type = ITEM_TRANSFORM;
    }
    else if (itemType == ITEM_POLYMESH)
    {
        tentry->matrix = curMatrix;
        tentry->type = ITEM_POLYMESH;

        if (widgetData->filteredShapes.size() > 0)
        {
            set<string>::iterator it =
                widgetData->filteredShapes.find(fullpath);
            if (it!=widgetData->filteredShapes.end())
            {
                processMesh(iObj, matrix, Entity, path);
            }
        }
        else
        {
            processMesh(iObj, matrix, Entity, path);
        }
    }
    else
    {
        tentry->matrix = curMatrix;
        tentry->type = ITEM_UNKNOWN;
    }

    M44d newMatrix = curMatrix * matrix;

    tree.children.push_back(tentry);

    for ( unsigned int i = 0 ; i < iObj.getNumChildren() ; i++ )
    {
        processObject(IObject(iObj, iObj.getChildHeader(i).getName()),
            newMatrix, Entity, *tentry);
    }

    return 1;
}


bool AlembicReader::processObjectData(MriGeoEntityHandle Entity)
{
    if (archive.valid())
    {
        root = archive.getTop();

        M44d matrix;
        matrix.makeIdentity();

        // start with the children of the root node
        for ( unsigned int i = 0 ; i < root.getNumChildren() ; i++ )
        {
            processObject(IObject(root, root.getChildHeader(i).getName()),
                matrix, Entity, widgetData->rootTree);
        }

    }
    else
        return false;

    return true;
}


bool AlembicReader::loadAlembic(MriGeoEntityHandle Entity)
{

    bool result = loadArchive(pFileName);
    if (!result)
    {
        cout << "Alembic archive load failed." << endl;
        return false;
    }

    // grab data from alembic file and ultimately create geometry in mari
    processObjectData(Entity);

    freeArchive();

    return true;
}


}
